//*********************************************************
//
//    Copyright (c) Microsoft. All rights reserved.
//    This code is licensed under the Apache License, Version 2.0.
//    THIS CODE IS PROVIDED *AS IS* WITHOUT WARRANTY OF
//    ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING ANY
//    IMPLIED WARRANTIES OF FITNESS FOR A PARTICULAR
//    PURPOSE, MERCHANTABILITY, OR NON-INFRINGEMENT.
//
//*********************************************************

using System;
using System.Collections.Generic;
using System.Text;
using ER=Microsoft.Research.eResearch;
using SH = Microsoft.Hpc.Scheduler;
using SHP = Microsoft.Hpc.Scheduler.Properties;

namespace Trident.Scheduler.Hpc
{
    class HPCSchedulerWrapper 
    {
        private string headNodeName = string.Empty;
        private SH.IScheduler scheduler = new SH.Scheduler();
        
        /// <summary>
        /// Field name should be used with HPC custom data
        /// </summary>
        public static string tridentJobIDFieldName = @"TridentJobID";


        /// <summary>
        /// HPC NodeGroup name for compute nodes
        /// </summary>
        const string computeNodeGroup = @"ComputeNodes";

        /// <summary>
        /// ComputeNodeList
        /// </summary>
        public List<string> ComputeNodeList
        {
            get { return this.computeNodeList; }            
        }

        /// <summary>
        /// Compute node list
        /// </summary>
        private List<string> computeNodeList = new List<string>();


        /// <summary>
        /// workflow laoder runtime  path to launch on other system
        /// </summary>
        private string tridentLauncherPath;


        /// <summary>
        /// Delegate to define the status change
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        public delegate void TaskStatusChanged(object sender, SH.TaskStateEventArg e);

        /// <summary>
        /// Task status changed event
        /// </summary>
        public event TaskStatusChanged OnTaskStatusChanged;

        /// <summary>
        /// Delegate to define the status change
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="e"></param>
        public delegate void JobStatusChanged(object sender, SH.JobStateEventArg e);

        /// <summary>
        /// Task status changed event
        /// </summary>
        public event JobStatusChanged OnJobStatusChanged;

        /// <summary>
        /// This method returns the current scheduler
        /// </summary>
        public SH.IScheduler Scheduler
        {
            get 
            {
                return scheduler; 
            }
        }
        /// <summary>
        /// Property to manage the launcher path
        /// </summary>
        public string TridentLauncherPath
        {
            get 
            {
                return tridentLauncherPath;
            }
            set 
            {
                tridentLauncherPath = value;
            }
        }
        /// <summary>
        /// Constructor to initlize wrapper object
        /// </summary>
        public HPCSchedulerWrapper()
        {
            
        }
        /// <summary>
        /// Initializes the class and connects to head node
        /// </summary>
        /// <param name="nodeName"></param>
        public void Initialize(string nodeName, string runtimePath)
        {
            tridentLauncherPath = runtimePath;

            ConnectToHeadNode(nodeName);

            //Add the compute node list
            SH.IStringCollection nodesList = scheduler.GetNodesInNodeGroup(computeNodeGroup);
            foreach (string node in nodesList)
            {
                computeNodeList.Add(node);
            }
           
        }

        private void ConnectToHeadNode(string nodeName)
        {
            headNodeName = Environment.MachineName;
            if (!String.IsNullOrEmpty(nodeName))
            {
                headNodeName = nodeName;
            }
            scheduler.Connect(headNodeName);

        }

        /// <summary>
        /// This method returns the number of cores for the node, failure returns 0
        /// </summary>
        /// <param name="nodeName"></param>
        /// <returns></returns>
        public int GetNumberOfCores(string nodeName)
        {
            int coresPerNode = 0;
            foreach (SH.ISchedulerNode node in scheduler.GetNodeList(null, null))
            {
                if (node.Name.Equals(nodeName, StringComparison.OrdinalIgnoreCase))
                {
                    coresPerNode = node.NumberOfCores;
                    break;
                }
            }
            return coresPerNode;
        }

        /// <summary>
        /// This method calculates the cores which can be allocated per
        /// job. The formula for doing this is
        /// (No Of Cores) = (Num of Cores in the selected node) / (Max concurrent Jobs configured on the Node)
        /// </summary>
        /// <param name="nodeName"></param>
        /// <returns></returns>
        private int GetCoresPerJob(string nodeName, int maxConcurrentJobs)
        {
            int coresPerJob = 1;
            foreach (SH.ISchedulerNode node in scheduler.GetNodeList(null, null))
            {
                if (node.Name.Equals(nodeName, StringComparison.OrdinalIgnoreCase))
                {
                    int numCores = node.NumberOfCores;
                    if (0 < maxConcurrentJobs)
                    {
                        coresPerJob = numCores / maxConcurrentJobs;
                    }
                    break;
                }
            }
            return coresPerJob;
        }

        /// <summary>
        /// Creates and HPC job with properties passed on arguments to this method 
        /// </summary>
        /// <param name="runATNodes"></param>
        /// <param name="TridentJobGuid"></param>
        /// <param name="ProjectName"></param>
        /// <param name="jobName"></param>
        /// <param name="priority"></param>
        /// <param name="exclusive"></param>
        /// <param name="maxConcurrentJobs"></param>
        /// <returns></returns>
        public SH.ISchedulerJob CreateJob(string[] runATNodes, 
                                        string TridentJobGuid, 
                                        string ProjectName, 
                                        string jobName,
                                        SHP.JobPriority priority,
                                        bool exclusive,
                                        int maxConcurrentJobs)                                        
        {
            #region param validation
            //Check whether the node is specified to run.
            if (runATNodes.Length <= 0 )
            {
                throw new Exception("Nodes are not specified for creating new job");
            }

            if (string.IsNullOrEmpty(TridentJobGuid))
            {
                throw new Exception("Required Trident Job ID missing, to schedule a job from trident");
            }

            #endregion

           
            #region create job
            //Create Job
            SH.ISchedulerJob hpcJob = scheduler.CreateJob();

            hpcJob.AutoCalculateMax = false;
            hpcJob.AutoCalculateMin = false;
            hpcJob.Name = jobName;

            #endregion

            #region update hpcJob properties
            //Update the requested nodes
            SH.IStringCollection requestedNodes = new SH.StringCollection();
            foreach (string node in runATNodes)
            {
                requestedNodes.Add(node);                
            }
            hpcJob.RequestedNodes = requestedNodes;
            
            //Change the priority based on hpcJob priority
            hpcJob.Priority = priority;
            hpcJob.Project = ProjectName;
            hpcJob.IsExclusive = exclusive;
            hpcJob.SetCustomProperty(tridentJobIDFieldName, TridentJobGuid);

            //If HPC selects the node it might not be accurate, selecting the head node processor count
            int processorCount = GetCoresPerJob(runATNodes[0], maxConcurrentJobs);
            hpcJob.MinimumNumberOfCores = processorCount;
            hpcJob.MaximumNumberOfCores = processorCount;

            //Register for task state events
            hpcJob.OnJobState += new EventHandler<SH.JobStateEventArg>(OnJobState);
            hpcJob.OnTaskState += new EventHandler<SH.TaskStateEventArg>(OnTaskState);

            #endregion update hpcJob properties

            return hpcJob;

        }


        /// <summary>
        /// Method to create a Hpc task on a specified Hpc job
        /// </summary>
        /// <param name="TridentJobGuid"></param>
        /// <param name="runATNodes"></param>
        /// <param name="ProjectName"></param>
        /// <param name="WorkflowName"></param>
        /// <param name="hpcJob"></param>
        /// <param name="hpcTaskName"></param>
        /// <returns></returns>
        public SH.ISchedulerTask CreateTask(string TridentJobGuid, 
                                            string[] runATNodes, 
                                            string taskName,
                                            SH.ISchedulerJob hpcJob)                                        
        {
            #region create Task
            //Create the command for Task 
            StringBuilder taskCommand = new StringBuilder();
            taskCommand.Append("\"");
            taskCommand.Append(this.tridentLauncherPath);
            taskCommand.Append("\" -j \"");
            taskCommand.Append(TridentJobGuid);
            taskCommand.Append("\"");

            //Create the new task 
            SH.ISchedulerTask task = hpcJob.CreateTask();
            task.Name = taskName;            

            task.CommandLine = taskCommand.ToString();
            task.MinimumNumberOfCores = hpcJob.MinimumNumberOfCores;
            task.MaximumNumberOfCores = hpcJob.MinimumNumberOfCores;
            #endregion
            //Add the task to job
            hpcJob.AddTask(task);

            return task;

        }


        /// <summary>
        /// Event for Task state change notification from HPC
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="args"></param>
        void OnTaskState(object sender, SH.TaskStateEventArg args)
        {
            if (null != OnTaskStatusChanged)
            {
                OnTaskStatusChanged(sender, args);
            }
        }

        /// <summary>
        /// Event for job state change notification from HPC
        /// </summary>
        /// <param name="sender"></param>
        /// <param name="args"></param>
        void OnJobState(object sender, SH.JobStateEventArg args)
        {
            if (null != OnJobStatusChanged)
            {
                OnJobStatusChanged(sender, args);
            }
        }
       
        
        /// <summary>
        /// Return the status of the job specified in the int jobid
        /// </summary>
        /// <param name="jobID"></param>
        /// <returns></returns>
        public SHP.JobState GetJobState(int jobID)
        {
            SH.ISchedulerJob job = scheduler.OpenJob(jobID);
            return job.State;            
        }

        /// <summary>
        /// Gets the task state of the sepcified job and task id
        /// </summary>
        /// <param name="jobID"></param>
        /// <param name="jobTaskID"></param>
        /// <returns></returns>
        public SHP.TaskState GetTaskState(int jobID, string taskName)
        {
            SHP.TaskState state = SHP.TaskState.Failed;
            SH.ISchedulerJob job = scheduler.OpenJob(jobID);

            SH.IFilterCollection filters = null;
            filters = scheduler.CreateFilterCollection();
            filters.Add(SHP.FilterOperator.Equal, SH.PropId.Task_Name, taskName);

            SH.ISchedulerCollection taskCollection = job.GetTaskList(filters, null, true);
            foreach (SH.ISchedulerTask task in taskCollection)
            {
                state = task.State;
                break;
            }            
            return state;
        }        
    }
}